import torch
from torchvision import transforms
from PIL import Image
import model  # 确保这个模块是你的 model.py 文件

# 类别名称
class_names = ['harmful', 'kitchen', 'other', 'recyclable']

# 加载模型
def load_model(model_path):
    # 初始化模型
    model_instance = model.initialize_model(num_classes=4)
    # 加载模型权重
    model_instance.load_state_dict(torch.load(model_path, map_location='cpu'))
    # 设置为评估模式
    model_instance.eval()
    return model_instance

# 预测函数
def predict_image(image_path, model_instance):
    # 预处理
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    # 加载并预处理图像
    image = Image.open(image_path).convert("RGB")  # 确保图像是 RGB 格式
    image = transform(image).unsqueeze(0)  # 添加批次维度
    
    # 预测
    with torch.no_grad():
        outputs = model_instance(image)
        _, predicted = torch.max(outputs, 1)
        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
    
    # 获取结果
    predicted_class = class_names[predicted[0]]
    confidence = probabilities[predicted[0]].item()
    
    # 获取所有类别的概率
    all_probs = {class_names[i]: probabilities[i].item() for i in range(len(class_names))}
    
    return predicted_class, confidence, all_probs